from alembic import op
import sqlalchemy as sa

def upgrade():
    # 扩展现有模型表
    op.add_column('model', sa.Column('parent_model_id', sa.Integer(), nullable=True))
    op.add_column('model', sa.Column('model_type', sa.String(50), default='base'))
    op.add_column('model', sa.Column('tags', sa.Text(), nullable=True))  # 移除默认值
    op.add_column('model', sa.Column('model_size', sa.BigInteger(), nullable=True))
    op.add_column('model', sa.Column('training_dataset', sa.String(500), nullable=True))
    op.add_column('model', sa.Column('validation_metrics', sa.Text(), nullable=True))  # 移除默认值
    op.add_column('model', sa.Column('model_config', sa.Text(), nullable=True))  # 移除默认值
    op.add_column('model', sa.Column('dependencies', sa.Text(), nullable=True))  # 移除默认值
    
    # 创建版本管理表
    op.create_table('model_version',
        sa.Column('id', sa.Integer(), primary_key=True),
        sa.Column('model_id', sa.Integer(), sa.ForeignKey('model.id')),
        sa.Column('version_number', sa.String(50)),
        sa.Column('version_type', sa.String(20)),
        sa.Column('changelog', sa.Text()),
        sa.Column('is_active', sa.Boolean(), default=True),
        sa.Column('performance_diff', sa.Text(), nullable=True),  # 移除默认值
        sa.Column('created_on', sa.DateTime()),
        sa.Column('changed_on', sa.DateTime())
    )
    
    # 创建依赖关系表
    op.create_table('model_dependency',
        sa.Column('id', sa.Integer(), primary_key=True),
        sa.Column('source_model_id', sa.Integer(), sa.ForeignKey('model.id')),
        sa.Column('target_model_id', sa.Integer(), sa.ForeignKey('model.id')),
        sa.Column('dependency_type', sa.String(50)),
        sa.Column('dependency_strength', sa.Float(), default=1.0),
        sa.Column('description', sa.String(500))
    )